import torch
from torch import nn
from torchvision import transforms
from Module import Utils as myutils

# 定义模型
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.structure = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2),
            nn.Sigmoid(),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0),
            nn.Sigmoid(),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Flatten(),
            nn.Linear(in_features=16*5*5, out_features=120),
            nn.Sigmoid(),
            nn.Linear(in_features=120, out_features=84),
            nn.Sigmoid(),
            nn.Linear(in_features=84, out_features=10)
        )

    def forward(self, x):
        return self.structure(x)
    
# 图片预处理
trans = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Resize(size=(28,28))
    ]
)

# 设置超参数
batch_size = 128
epochs = 20
lr = 1.2

# 获取数据集
_, train_iter, test_iter = myutils.load_data_FMNIST('../data', batch_size, trans)

# 实例化网络
net = LeNet()

# 创建优化器
optimizer = torch.optim.SGD(params=net.parameters(),lr=lr)

# 定义损失函数
loss = nn.CrossEntropyLoss()

# 开始训练
myutils.train(net, train_iter, test_iter, epochs, optimizer, loss, True)